import os

import fitz  # PyMuPDF
import pandas as pd
from tqdm import tqdm

save_dir = 'data'

keyword_li = [
    '生成式', '大模型', 'transformer', 'gpt', 'bert', 'BERT', 'GPT', 'Transformer',
    'GAN'
]


def search_keyword_in_pdf(pdf_path, keywords):
    doc = fitz.open(pdf_path)
    pg_data = []
    data = {}
    for page_num in range(len(doc)):
        page = doc[page_num]
        # print(dir(page))
        text = page.getText('text')
        for k in keywords:
            if k in text:
                pg_data.append(page_num + 1)
                if k in data:
                    data[k].append(page_num + 1)
                else:
                    data[k] = [page_num + 1]

                # data += f'"{k}" 出现在第 {page_num + 1} 页；'
    return pg_data, '；'.join([f'"{k}" 出现在第 {v} 页' for k, v in data.items()])


ret_data = []

for filename in tqdm(os.listdir(save_dir)):
    # print(filename)
    save_path = os.path.join(save_dir, filename)
    try:
        page_num_li, result = search_keyword_in_pdf(save_path, keyword_li)
    except Exception as e:
        print(e)
        print(save_path)
        continue
    if page_num_li:
        ret_data.append({
            'name': filename,
            'result': result,
            'page': page_num_li
        })

pd_data = pd.DataFrame(ret_data)

print(pd_data)

pd_data.to_csv('企业年报解析结果.csv', index=None)
